import math
from locust import LoadTestShape
from utils.time_tool.time_control import parse_duration


# === 统一的基类，支持配置注入 ===
class ConfigurableLoadShape(LoadTestShape):
    def __init__(self, users=1, spawn_rate=1, duration="1m", stages=None):
        super().__init__()
        self.shape_users = users
        self.spawn_rate = spawn_rate
        self.shape_duration = parse_duration(duration)
        self.shape_stages = stages or []

    def tick(self):
        print(f"负载控制器 tick() 调用: 当前用户数 = {self.shape_users}, spawn_rate = {self.spawn_rate}")
        if self.get_run_time() < self.shape_duration:
            return (self.shape_users, self.spawn_rate)
        return None




# === 各种负载模式实现 ===

class BaselineShape(ConfigurableLoadShape):
    """恒定负载控制器"""
    def tick(self):
        print("+++++++++++走到了恒定负载的控制器这里+++++++++++++++")
        print(f"负载控制器 tick() 调用: 当前用户数 = {self.shape_users}, spawn_rate = {self.spawn_rate}")
        if self.get_run_time() < self.shape_duration:
            return (self.shape_users, self.spawn_rate)
        return None


class IncrementalLoadShape(ConfigurableLoadShape):
    """阶梯递增负载控制器"""
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.current_stage = 0

    def tick(self):
        run_time = self.get_run_time()
        stages = self.shape_stages

        while self.current_stage < len(stages):
            stage = stages[self.current_stage]
            if run_time < stage["duration"]:
                return (stage["users"], stage["spawn_rate"])
            run_time -= stage["duration"]
            self.current_stage += 1
        return None


class StressTestShape(ConfigurableLoadShape):
    """峰值压力控制器"""
    def tick(self):
        if self.get_run_time() < self.shape_duration:
            return (self.shape_users, self.spawn_rate)
        return None


class ConcurrencyShape(ConfigurableLoadShape):
    """瞬时并发控制器"""
    def tick(self):
        if self.get_run_time() < self.shape_duration:
            return (self.shape_users, self.spawn_rate)
        return None


class WaveShape(ConfigurableLoadShape):
    """波浪型负载控制器"""
    def tick(self):
        current_time = self.get_run_time()
        if current_time < self.shape_duration:
            users = int(100 + 50 * math.sin(current_time / 60))
            return (users, self.spawn_rate)
        return None


# === 负载模式映射 ===
SHAPE_CLASS_MAP = {
    "constant": BaselineShape,         # 恒定
    "incremental": IncrementalLoadShape,  # 阶梯递增
    "peak": StressTestShape,           # 峰值
    "concurrency": ConcurrencyShape,   # 并发瞬发
    "wave": WaveShape                  # 波动型
}


# === 工厂函数：根据配置创建 LoadShape 实例 ===
def create_load_shape(profile_config):
    print("最终传入的配置是：",profile_config)
    shape_type = profile_config.get("load_shape", "constant")
    shape_class = SHAPE_CLASS_MAP.get(shape_type, BaselineShape)
    return shape_class(
        users=profile_config.get("users", 2),
        spawn_rate=profile_config.get("spawn_rate", 1),
        duration=profile_config.get("duration", "1m"),
        stages=profile_config.get("stages", [])
    )


# class DebugShape(LoadTestShape):
#     def tick(self):
#         print("✅ DebugShape tick 调用了！！！")
#         return (10, 2)